"""
RFB protocol implementation, client side

References:
- http://www.realvnc.com/docs/rfbproto.pdf
- https://github.com/rfbproto/rfbproto/blob/master/rfbproto.rst
"""
import base64
import datetime
import io
import json
import struct
import pyDes
import urllib3
from PIL import Image
from threading import Thread
import logging
import socket
from socket import SHUT_RDWR
import struct as s
import time
import pyghmi.ipmi.command as ipmi_command
from pythonping import ping

urllib3.disable_warnings()
HTTP = urllib3.PoolManager(cert_reqs='CERT_NONE')

logger = logging.getLogger("vnc_screenshot")
logging.basicConfig(format="[%(name)s] %(levelname)s: %(message)s", level=logging.INFO)

## Encoding Type for SetEncodings()
# publicly documented
ENC_RAW = 0  # Raw
ENC_COPYRECT = 1  # CopyRect
ENC_RRE = 2  # RRE
ENC_HEXTILE = 5  # Hextile
ENC_TRLE = 15  # TRLE
ENC_ZRLE = 16  # ZRLE
ENC_JPEG = 21  # JPEG Encoding

# Authentication protocol types
AUTH_FAIL = 0
AUTH_NONE = 1
AUTH_VNCAUTH = 2

# Authentication result types
SMSG_AUTH_OK = 0
SMSG_AUTH_FAIL = 1
SMSG_AUTH_TOOMANY = 2

# Server message types
SMSG_FBUPDATE = 0
SMSG_SETCOLORMAP = 1
SMSG_BELL = 2
SMSG_SERVERCUTTEXT = 3

# Client message types
CMSG_SETPIXELFORMAT = 0
CMSG_SETENCODINGS = 2
CMSG_FBUPDATEREQ = 3
CMSG_KEYEVENT = 4
CMSG_POINTEREVENT = 5
CMSG_CLIENTCUTTEXT = 6

# pseudo-encodings
ENC_CURSOR = -239  # Cursor position pseudo-encoding
ENC_DESKTOPSIZE = -223  # DesktopSize pseudo-encoding

MAX_BUFF_SIZE: int = 10 * 1024 * 1024  # 10MB

SUPPORTED_VERSIONS = [
    (3, 8)
]
KNOWN_VERSIONS = [
    (3, 3), (3, 6), (3, 7), (3, 8),
    (4, 0), (4, 1),
    (5, 0)
]

SUPPORTED_ENCODINGS = [
    ENC_RAW
]


def return_uint8_bytes(value: int, big_endian=False) -> bytes:
    if big_endian:
        return struct.pack(">B", value)
    else:
        return struct.pack("<B", value)


def return_uint8_val(data: bytes, big_endian=False) -> int:
    if big_endian:
        return struct.unpack(">B", data)[0]
    else:
        return struct.unpack("<B", data)[0]


def return_sint32_bytes(value: int, big_endian=False) -> bytes:
    if big_endian:
        return struct.pack(">i", value)
    else:
        return struct.pack("<i", value)


def write_sint32_buff(buffer, value: int, big_endian=False) -> None:
    buffer.write(return_sint32_bytes(value, big_endian))


class RestfulClient:
    def __init__(self, host, username, password):
        self.host = host
        self.username = username
        self.password = password

        self.base_url = f"https://{self.host}"
        self.headers = {'Content-Type': 'application/json'}

    def session(self):
        """创建连接会话，获取token"""

        url = self.base_url + "/api/session"
        fields = {
            "username": self.username,
            "password": self.password,
            "log_Type": "1",  # 加密
        }

        resp = HTTP.request("POST", url, retries=False, fields=fields, timeout=10)
        result = json.loads(resp.data.decode())

        if isinstance(result, dict) and result.get("cc", None) == 0:
            try:
                self.headers['X-CSRFTOKEN'] = result["CSRFToken"]
                # 通过关键字截取字段获取cookie值
                cookie = resp.headers.get("set-cookie")
                session_id = str(cookie).split("ID=")[-1].split("; Pa")[0]
                self.headers['Cookie'] = "QSESSIONID=%s;CSRF=%s" % (session_id, result["CSRFToken"])
            except (IndexError, KeyError, Exception) as err:
                raise ValueError(err)
        else:
            err_message = ("Failure: failed to establish a new "
                           "connection to the host: response abnormal")
            raise ValueError(err_message)

    def request(self, method, uri, body=None):
        """封装的请求函数"""
        if not uri.startswith('/'):
            uri = '/' + uri

        url = self.base_url + uri
        if body:
            payload = json.dumps(body)
            resp = HTTP.request(method, url, body=payload,
                                retries=False, headers=self.headers, timeout=20)
        else:
            resp = HTTP.request(method, url, retries=False, headers=self.headers, timeout=20)

        result = json.loads(resp.data.decode())
        if resp.status != 200 and result.get("cc") == 0:
            raise ValueError(f'request err: {resp.reason}')

        return result

    def delete_session(self):
        """释放会话资源"""
        self.request("DELETE", "/api/session")

    def get_storage(self):
        url = '/api/system/storageinfo/hdd/list'
        resp = self.request("get", url)
        dis_phys_info = resp['dis_phys_info']
        _error = []
        for dis_phys_item in dis_phys_info:
            if dis_phys_item['health_state'] > 0:
                _error.append(f"硬盘故障 {dis_phys_item['fault_desc']} 序列号{dis_phys_item['serial']}")
        return _error

    def status_power(self):
        url = '/api/system/power_info'
        resp = self.request("get", url)
        # dis_phys_info = resp['dis_phys_info']
        _error = []
        for power_item in resp['detail']:
            if power_item['health_status'] > 0:
                _error.append(f"电源故障 {power_item['power_sn']} {power_item['fault_description']}")
        return _error

    def check_health(self):
        datas = []
        url = "/api/system_health"
        resp = self.request("get", url)
        if resp['overall_health']['health'] == 0:
            return datas
        fault_parts = []
        del resp['overall_health']
        del resp['cc']

        for k, v in resp.items():
            if v['health'] > 0:
                fault_parts.append(k)
        func = {
            'storage': self.get_storage,
            'psu': self.status_power,
        }
        for parts in fault_parts:
            if parts not in func:
                datas.append(f'未知故障：{parts}')
                continue
            datas.extend(func[parts]())
        return datas


class RFBPixelformat:
    def __init__(self,
                 bpp=32, depth=24, bigendian=False, truecolor=True,
                 redmax=255, greenmax=255, bluemax=255,
                 redshift=0, greenshift=0, blueshift=16):
        self.bitspp = bpp
        self.depth = depth
        self.bigendian = 1 if bigendian else 0
        self.truecolor = 1 if truecolor else 0

        self.redmax = redmax
        self.greenmax = greenmax
        self.bluemax = bluemax

        self.redshift = redshift
        self.greenshift = greenshift
        self.blueshift = blueshift

    @staticmethod
    def getRGB32():
        return RFBPixelformat(
            bpp=32, depth=32,
            redshift=16, greenshift=8, blueshift=0
        )

    @staticmethod
    def getRGB24():
        return RFBPixelformat(
            bpp=32, depth=24,
            redshift=16, greenshift=8, blueshift=0
        )

    @staticmethod
    def getRGB16():
        return RFBPixelformat(
            bpp=16, depth=16,
            redmax=31, greenmax=63, bluemax=31,
            redshift=11, greenshift=5, blueshift=0
        )

    @staticmethod
    def getRGB555():
        return RFBPixelformat(
            bpp=16, depth=15,
            redmax=31, greenmax=31, bluemax=31,
            redshift=10, greenshift=5, blueshift=0
        )

    @staticmethod
    def getRGB444():
        return RFBPixelformat(
            bpp=16, depth=12,
            redmax=15, greenmax=15, bluemax=15,
            redshift=8, greenshift=4, blueshift=0
        )

    def asTuple(self) -> tuple:
        return (
            self.bitspp, self.depth, self.bigendian, self.truecolor,
            self.redmax, self.greenmax, self.bluemax,
            self.redshift, self.greenshift, self.blueshift
        )

    def __str__(self) -> str:
        return ";".join(str(x) for x in self.asTuple())


class RFBDes(pyDes.des):
    def setKey(self, key):
        """
        RFB protocol for authentication requires client to encrypt
        challenge sent by server with password using DES method. However,
        bits in each byte of the password are put in reverse order before
        using it as encryption key.
        """
        newkey = list()
        for bsrc in key:
            btgt = 0
            for i in range(8):
                if bsrc & (1 << i):
                    btgt = btgt | (1 << 7 - i)
            newkey.append(btgt)
        super(RFBDes, self).setKey(newkey)


def des_encrypt_password(challenge, password):
    # 将密码调整为 8 字节
    password = password[:8]  # 截断到8字节
    password = password.ljust(8, '\0')  # 用零填充至8字节
    password = password.encode('utf-8')

    des = RFBDes(password)
    return des.encrypt(challenge)


class RFBClient:
    pixformat: RFBPixelformat
    numRectangles = 0

    _stop = False
    connected = None
    _requestFrameBufferUpdate = False
    _incrementalFrameBufferUpdate = True

    def __init__(self, host, port=5900, password: str = None):
        self.host = host
        self.port = port
        self.password = password
        self._mainLoop: Thread = None
        self.screen_refresh_count = 0
        self.screen_images = []

    def __start(self):
        """建立vnc连接，客户端客户使用3.8版本"""
        print(self.host, self.port)
        self.connection = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        try:
            self.connection.connect((self.host, self.port))
        except ConnectionRefusedError:
            print(f'{self.host} 连接失败')
            self.connected = False
            return

        buffer = self.connection.recv(12, socket.MSG_WAITALL)

        # 确认RFB协议版本号
        if b'\n' in buffer and buffer.startswith(b'RFB'):
            _maj, _min = [int(x) for x in buffer[3:-1].split(b'.')]
        else:
            self.__close()
            raise Exception("RFB protocol version mismatch")

        logger.info(f'VNC server version: {_maj}.{_min}')
        logger.info(f'VNC client version: 3.8')

        self.version_maj, self.version_min = 3, 8

        # 协商安全类型
        self.connection.send(f"RFB 00{self.version_maj}.00{self.version_min}\n".encode())
        logger.info("Connected to VNC server")

        data = self.connection.recv(32)
        # 3.8版本 解码为无符号整数，按照大端字节序
        auth_types = struct.unpack(f"{len(data)}B", data)
        logger.info(f'Security types: {auth_types}')
        if not auth_types:
            raise ValueError('未获取到认证协议！')

        if auth_types[-1] == 2:
            """使用密码验证"""
            self.connection.send(b'\x02')
            # 服务器返回的16字节挑战数据
            challenge = self.connection.recv(16)

            # 加密挑战数据
            encrypted_challenge = des_encrypt_password(challenge, self.password)
            self.connection.send(encrypted_challenge)  # 发送加密的挑战数据
            data = self.connection.recv(4, socket.MSG_WAITALL)  # 接收验证响应
            auth_result = struct.unpack("<I", data)[0]
            if auth_result == SMSG_AUTH_OK:
                logger.info('VNC Authentication successful')
                shared = 1  # 1 表示共享，不中断别的连接，否则为0
                self.connection.send(return_uint8_bytes(shared, True))
                self._handle_server_init(self.connection.recv(24, socket.MSG_WAITALL))
            else:
                logger.error("VNC Authentication failed")

        elif auth_types[-1] == 1:
            """无需密码验证"""
            raise ValueError('不支持无密码认证！')
        elif auth_types[-1] == 0:
            raise ValueError('认证失败')
        else:
            self.__close()
            raise Exception(f'Unknown auth response {auth_types}')

        return True

    def __close(self):
        logger.info("Closing connection")

        if self.connection:
            try:
                self.connection.shutdown(SHUT_RDWR)
                self.connection.close()
            except OSError:
                logger.info("TCP Connection already closed")

    def _handle_server_init(self, data: bytes):
        try:
            self.vncWidth, self.vncHeight, pixformat, namelen = s.unpack("!HH16sI", data)
        except s.error as e:
            logger.error("Handshake failed")
            self.__close()
            raise ValueError(e)

        self.desktopname = self.connection.recv(namelen, socket.MSG_WAITALL).decode()
        logger.debug(f"Connecting to \"{self.desktopname}\"")

        pixformatData = s.unpack("!BBBBHHHBBBxxx", pixformat)
        self.pixformat = RFBPixelformat(*pixformatData)

        logger.debug(f"Server Pixelformat: {self.pixformat}")
        logger.debug(f"Resolution: {self.vncWidth}x{self.vncHeight}")

        self.set_encodings(SUPPORTED_ENCODINGS)

        self.onConnectionMade()
        self.connected = True

        self._main_request_loop()

    def onConnectionMade(self):
        self.setPixelFormat(RFBPixelformat.getRGB32())

    def _main_request_loop(self):
        time.sleep(0.5)
        self.framebuffer_update_request(incremental=False)

        while not self._stop and self.connection:
            d_type = self.connection.recv(1)
            if not d_type:
                continue
            self._handle_connection(d_type)

            self.framebuffer_update_request(incremental=False)
            time.sleep(0.5)

        logger.debug("loop exit")

    def _handle_connection(self, data: bytes):
        msg_id = return_uint8_val(data)
        if msg_id == SMSG_FBUPDATE:
            self._handle_framebuffer_update(self.connection.recv(3, socket.MSG_WAITALL))
        else:
            logger.warning(f"Unknown message type recieved (id {msg_id})")

    def _handle_framebuffer_update(self, data: bytes):
        """帧缓冲更新"""
        self._handle_rectangle(self.connection.recv(12, socket.MSG_WAITALL))

    def _handle_rectangle(self, data: bytes):
        """更新屏幕数据"""

        x, y, width, height, encoding = s.unpack("!HHHHi", data)
        logger.info(f"Encoding: {encoding}")

        if encoding == ENC_DESKTOPSIZE:
            logger.info(f"ENC_DESKTOPSIZE skipping")

        elif encoding == ENC_RAW:
            size = (width * height * self.pixformat.bitspp) // 8
            logger.debug(f"expected size: {size}")
            data = self.connection.recv(size, socket.MSG_WAITALL)
            logger.info(f"ENC_RAW skipping. length: {len(data)}")

        elif encoding == ENC_JPEG:
            jpeg_data = bytearray()
            while True:
                # 接收数据
                _data = self.connection.recv(1024)
                if not _data:
                    break
                jpeg_data.extend(_data)

                # 完整JPEG图像的结束标记： b'\xFF\xD9'
                if jpeg_data[-2:] == b'\xFF\xD9':
                    break

            if jpeg_data:
                self.screen_refresh_count += 1

                # 压缩数据
                image = Image.open(io.BytesIO(jpeg_data))
                _output = io.BytesIO()
                image.save(_output, quality=10, optimize=True, format="JPEG")

                self.screen_images.append(base64.b64encode(_output.getvalue()).decode('utf-8'))

        else:
            logger.error(f"Unsupported encoding received ({encoding}), skipping.")
            raise TypeError(f"Unsupported encoding received ({encoding})")

    def setPixelFormat(self, pixelformat: RFBPixelformat):
        logger.debug(f"Requesting pixelformat: {pixelformat}")

        self.pixformat = pixelformat
        pformat = s.pack("!BBBBHHHBBBxxx", *pixelformat.asTuple())
        self.connection.send(s.pack("!Bxxx16s", CMSG_SETPIXELFORMAT, pformat))

    def set_encodings(self, encodings: list):

        self.connection.send(s.pack("!BxH", CMSG_SETENCODINGS, len(encodings)))
        for encoding in encodings:
            self.connection.send(return_sint32_bytes(encoding, True))

    def framebuffer_update_request(self,
                                   xPos=0, yPos=0,
                                   width=None, height=None,
                                   incremental=False):

        if not width:
            width = self.vncWidth - xPos
        if not height:
            height = self.vncHeight - yPos

        inc = 1 if incremental else 0

        x_position = 0
        y_position = 0

        # 构造请求消息，按文档格式
        message = struct.pack("!BBHHHH", 3, inc, x_position, y_position, width, height)

        self.connection.send(message)

    def start_connection(self):
        # return self.__start()
        self._mainLoop = Thread(target=self.__start)
        self._mainLoop.start()

    def close_connection(self):
        self._stop = True
        self.connected = False
        self.__close()
        logger.info('Close connection')


def check_powser(host):
    ipmi_session = ipmi_command.Command(host, 'cmssuser', ipmi_password)
    response = ipmi_session.get_power()
    res_dict = {
        'on': '开机中',
        'off': '已关机'
    }
    return res_dict[response['powerstate']]


def check_ping(host):
    response = ping(host, count=1, timeout=2)  # 发送1次ping请求，设置超时时间为2秒
    if response.success():
        return '管理网正常'
    else:
        return '管理网异常！'


def check_health_h3c(host):
    rc = RestfulClient(host=host, username='cmssuser', password=ipmi_password)
    rc.session()
    data = rc.check_health()
    print('data', data)

    rc.delete_session()

    return data


if __name__ == '__main__':
    # brand = '#{brand}'
    # ipmi_ip = '#{ipmi_ip}'
    # ipmi_password = '#{ipmi_password}'
    # manage_ip = '#{manage_ip}'
    output = []

    powser_status = check_powser(ipmi_ip)
    output.append({'name': '电源状态', 'status': powser_status})

    ping_status = check_ping(manage_ip)
    output.append({'name': '管理网状态', 'status': ping_status})
    if brand == 'H3C':
        fault = check_health_h3c(ipmi_ip)
        if not fault:
            fault = '健康'
        output.append({'name': '硬件健康状态', 'status': fault})

    if powser_status == '开机中':
        port = 5900
        password = 'EOt77##f'
        try:
            rbf = RFBClient(ipmi_ip, port, password)
            rbf.start_connection()

            wait_index = 0
            wait_refresh = 0
            while wait_index < 60:
                time.sleep(1)
                wait_index += 1

                if rbf.connected:
                    if rbf.screen_refresh_count > 1:
                        output.append({
                            'name': '操作系统状态',
                            'status': '正常',
                            'image': rbf.screen_images[0]
                        })
                        break
                    elif rbf.screen_refresh_count == 1:
                        if wait_refresh > 10:
                            output.append({
                                'name': '操作系统状态',
                                'status': '卡死',
                                'image': rbf.screen_images[0]
                            })
                            break
                        wait_refresh += 1


                elif wait_index > 10:
                    print(f'{ipmi_ip} 连接超时')
                    break

                elif rbf.connected is False:
                    output.append({
                        'name': '操作系统状态',
                        'status': 'VNC连接失败，请检查是否开启',
                    })
                    break

            rbf.close_connection()
        except Exception as e:
            output.append({
                'name': '操作系统状态',
                'status': e,
            })
    print('----------json output-----------')
    print(json.dumps(output, ensure_ascii=False))

